// Copyright 2020 PAL Robotics S.L.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*
 * Author: Enrique Fernández
 */

#include "four_wheel_steering_controller/odometry.hpp"
#include <iostream>

namespace four_wheel_steering_controller {
    Odometry::Odometry(size_t velocity_rolling_window_size)
        : timestamp_(0.0),
        x_(0.0),
        y_(0.0),
        heading_(0.0),
        linear_(0.0),
        angular_(0.0),
        wheel_track_(0.0),
        wheel_base_(0.0),
        wheel_radius_(0.0),
        wheel_num_(0.0),
        // wheel_old_pos_({ 0.0 }),
        // steering_old_pos_({ 0.0 }),
        velocity_rolling_window_size_(velocity_rolling_window_size),
        linear_x_accumulator_(velocity_rolling_window_size),
        linear_y_accumulator_(velocity_rolling_window_size),
        angular_accumulator_(velocity_rolling_window_size) {
    }

    void Odometry::init(const rclcpp::Time& time) {
        // Reset accumulators and timestamp:
        resetAccumulators();
        timestamp_ = time;
    }

    bool Odometry::update(double(&wheel_pos)[4], double(&steering_vel)[4], const rclcpp::Time& /*time*/) {
        // We cannot estimate the speed with very small time intervals:
        // const double dt = time.seconds() - timestamp_.seconds();
        // if (dt < 0.0001) {
        //     return false;  // Interval too small to integrate with
        // }

        // // Get current wheel joint positions:
        // const double left_wheel_cur_pos = left_pos * left_wheel_radius_;
        // const double right_wheel_cur_pos = right_pos * right_wheel_radius_;

        // // Estimate velocity of wheels using old and current position:
        // const double left_wheel_est_vel = left_wheel_cur_pos - left_wheel_old_pos_;
        // const double right_wheel_est_vel = right_wheel_cur_pos - right_wheel_old_pos_;

        // // Update old position with current:
        // left_wheel_old_pos_ = left_wheel_cur_pos;
        // right_wheel_old_pos_ = right_wheel_cur_pos;

        // updateFromVelocity(left_wheel_est_vel, right_wheel_est_vel, time);

        return true;
    }

    bool Odometry::updateFromVelocity(double(&wheel_vel)[4], double(&steering_pos)[4], const rclcpp::Time& time) {
        const double dt = time.seconds() - timestamp_.seconds();

        double v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y;

        v1x = wheel_vel[0] * cos(steering_pos[0]);
        v1y = wheel_vel[0] * sin(steering_pos[0]);
        v2x = wheel_vel[1] * cos(steering_pos[1]);
        v2y = wheel_vel[1] * sin(steering_pos[1]);
        v3x = wheel_vel[2] * cos(steering_pos[2]);
        v3y = wheel_vel[2] * sin(steering_pos[2]);
        v4x = wheel_vel[3] * cos(steering_pos[3]);
        v4y = wheel_vel[3] * sin(steering_pos[3]);

        // std::cout << "Odometry::v1x: " << v1x << " v1y: " << v1y << " v2x: " << v2x << " v2y: " << v2y << " v3x: " << v3x << " v3y: "<< v3y << " v4x: " << v4x << " v4y: " << v4y << std::endl;

        if (wheel_num_ == 4) {
            const double linear_x = (v1x + v2x + v3x + v4x) / 4.0;
            const double linear_y = (v1y + v2y + v3y + v4y) / 4.0;
            const double angular1 = (v3x - v1x + v4x - v2x) / wheel_track_ / 2.0;
            const double angular2 = (v1y - v3y + v4y - v2y) / wheel_base_ / 2.0;
            // const double linear = std::hypot(linear_x, linear_y);
            const double angular = (angular1 + angular2) / 2.0;

            integrateExact(linear_x * dt, linear_y * dt, angular * dt);

            // printf("linear: %f, angular: %f\n", linear, angular);
            // std::cout << "linear_x: " << linear_x << "linear_y: " << linear_y << " angular: " << angular << std::endl;
            timestamp_ = time;

            // Estimate speeds using a rolling mean to filter them out:
            linear_x_accumulator_.accumulate(linear_x);
            linear_y_accumulator_.accumulate(linear_y);
            angular_accumulator_.accumulate(angular);

            linear_x_ = linear_x_accumulator_.getRollingMean();
            linear_y_ = linear_y_accumulator_.getRollingMean();
            angular_ = angular_accumulator_.getRollingMean();

        }
        else if (wheel_num_ == 2) {
            const double linear_x = (v1x + v3x) / 4.0;
            const double linear_y = (v1y + v3y) / 4.0;
            const double angular1 = (v3x - v1x) / wheel_track_;
            const double angular2 = (v1y - v3y) / wheel_base_;
            const double angular = (angular1 + angular2) / 2.0;
            integrateExact(linear_x * dt, linear_y * dt, angular * dt);

            // printf("linear: %f, angular: %f\n", linear, angular);
            // std::cout << "linear_x: " << linear_x << "linear_y: " << linear_y << " angular: " << angular << std::endl;
            timestamp_ = time;

            // Estimate speeds using a rolling mean to filter them out:
            linear_x_accumulator_.accumulate(linear_x);
            linear_y_accumulator_.accumulate(linear_y);
            angular_accumulator_.accumulate(angular);

            linear_x_ = linear_x_accumulator_.getRollingMean();
            linear_y_ = linear_y_accumulator_.getRollingMean();
            angular_ = angular_accumulator_.getRollingMean();

        }
        else
            return  false;
        return true;
    }

    void Odometry::updateOpenLoop(double linear, double angular, const rclcpp::Time& time) {
        /// Save last linear and angular velocity:
        linear_ = linear;
        angular_ = angular;

        /// Integrate odometry:
        const double dt = time.seconds() - timestamp_.seconds();
        timestamp_ = time;
        integrateExact(linear * dt, angular * dt);
    }

    void Odometry::resetOdometry() {
        x_ = 0.0;
        y_ = 0.0;
        heading_ = 0.0;
    }

    void Odometry::setWheelParams(
        double         wheel_base,
        double         wheel_track,
        double         wheel_radius,
        int wheel_num) {
        wheel_base_ = wheel_base;
        wheel_track_ = wheel_track;
        wheel_radius_ = wheel_radius;
        wheel_num_ = wheel_num;

    }

    void Odometry::setVelocityRollingWindowSize(size_t velocity_rolling_window_size) {
        velocity_rolling_window_size_ = velocity_rolling_window_size;
        resetAccumulators();
    }

    void Odometry::integrateRungeKutta2(double linear, double angular) {
        const double direction = heading_ + angular * 0.5;

        /// Runge-Kutta 2nd order integration:
        x_ += linear * cos(direction);
        y_ += linear * sin(direction);
        heading_ += angular;
    }
    void Odometry::integrateRungeKutta2(double linear_x, double linear_y, double angular) {
        const double direction = heading_ + angular * 0.5;

        /// Runge-Kutta 2nd order integration:
        // x_ += linear * cos(direction);
        // y_ += linear * sin(direction);

        const  double x1 = linear_x * cos(direction);
        const double y1 = linear_x * sin(direction);
        const double x2 = linear_y * cos(direction + M_PI / 2);
        const double y2 = linear_y * sin(direction + M_PI / 2);

        x_ += (x1 + x2);
        y_ += (y1 + y2);

        heading_ += angular;
    }

    void Odometry::integrateExact(double linear, double angular) {
        if (fabs(angular) < 1e-6) {
            integrateRungeKutta2(linear, angular);
        }
        else {
            /// Exact integration (should solve problems when angular is zero):
            const double heading_old = heading_;
            const double r = linear / angular;
            heading_ += angular;
            x_ += r * (sin(heading_) - sin(heading_old));
            y_ += -r * (cos(heading_) - cos(heading_old));
        }
    }
    void Odometry::integrateExact(double linear_x, double linear_y, double angular) {

        /// Exact integration (should solve problems when angular is zero):
        // const double heading_old = heading_;
        // const double r = linear / angular;
        // x_ += linear_x;
        // y_ += linear_y;
        // x_ += linear_x * std::cos(heading_ + angular / 2) - linear_y * std::cos(heading_ + angular / 2+M_PI/2);
        // y_ += linear_x * std::sin(heading_ + angular / 2) + linear_y * std::sin(heading_ + angular / 2+M_PI/2);

        // heading_ += angular;
        // std::cout << "heading_: " << heading_ << "x_"<< x_ << "y_" << y_ << std::endl;



        if (fabs(angular) < 1e-6) {
            integrateRungeKutta2(linear_x, linear_y, angular);
        }
        else {
            const double heading_old = heading_;
            const double rx = linear_x / angular;
            const double ry = linear_y / angular;
            heading_ += angular;
            // x_ += r * (sin(heading_) - sin(heading_old));
            // y_ += -r * (cos(heading_) - cos(heading_old));

            double x1 = rx * (sin(heading_) - sin(heading_old));
            double y1 = -rx * (cos(heading_) - cos(heading_old));
            double x2 = ry * (sin(heading_ + M_PI / 2) - sin(heading_old + M_PI / 2));
            double y2 = -ry * (cos(heading_ + M_PI / 2) - cos(heading_old + M_PI / 2));

            x_ += (x1 + x2);
            y_ += (y1 + y2);
        }
        // std::cout << "heading_: " << heading_ << "  x_" << x_ << "  y_" << y_ << std::endl;

    }
    void Odometry::resetAccumulators() {
        linear_x_accumulator_ = RollingMeanAccumulator(velocity_rolling_window_size_);
        linear_y_accumulator_ = RollingMeanAccumulator(velocity_rolling_window_size_);
        angular_accumulator_ = RollingMeanAccumulator(velocity_rolling_window_size_);
    }

}  // namespace four_wheel_steering_controller
